#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 crane <crane@sharp>
#
# Distributed under terms of the MIT license.

"""

"""

import random


class Heap:
    def min_lessor(a, b):
        return a < b

    def max_lessor(a, b):
        return a > b

    def build_min(self, nums):
        # def lessor(a, b):
        #     return a < b
        return self.build(nums, Heap.min_lessor)

    def build_max(self, nums):

        return self.build(nums, Heap.max_lessor)

    def build(self, nums, lessor):
        # 通用build: 顶部放置最小元素
        start = (len(nums) - 2) // 2
        end = 0

        # print(start, end)
        for idx in range(start, end-1, -1):
            self.adjust(nums, idx, lessor)

        return nums

    # def insert_max(self, nums, value):
    #     lessor = Heap.max_lessor
    #     if lessor(value, nums[0]):
    #         nums[0] = value
    #         self.adjust(nums, 0, lessor)

    def adjust_min(self, nums, idx):
        self.adjust(nums, idx, Heap.min_lessor)

    def adjust_max(self, nums, idx, lessor):
        pass

    def adjust(self, nums, idx, lessor):
        # print("adjust", idx)
        if idx > (len(nums) - 2) // 2:
            return

        cur = nums[idx]
        left_idx  = idx * 2 + 1
        right_idx = idx * 2 + 2

        assert left_idx < len(nums)

        bottom_min_idx = 0
        if right_idx >= len(nums):
            print("right_idx greator", idx, len(nums))
            bottom_min_idx = left_idx
        else:
            if lessor(nums[left_idx], nums[right_idx]):
                bottom_min_idx = left_idx
            else:
                bottom_min_idx = right_idx

        if lessor(nums[bottom_min_idx], nums[idx]):
            nums[idx], nums[bottom_min_idx] = nums[bottom_min_idx], nums[idx]
            self.adjust(nums, bottom_min_idx, lessor)


class Solution:
    def __init__(self):
        pass

    def top_k(self, nums, k):
        if len(nums) < k:
            return nums

        h = Heap()
        heap_nums = h.build_min(nums[:k])

        # print(nums)
        # print(nums[k:])
        for e in nums[k:]:
            if heap_nums[0] < e:
                heap_nums[0] = e
                h.adjust_min(heap_nums, 0)

        return heap_nums


def test():
    s = Solution()

    # l = [1, 2, 3, 4]
    # s.top_k(l, 4)

    for max_num in range(3, 998):
        l = list(range(max_num))
        random.shuffle(l)
        k = 7
        tops = s.top_k(l, k)

        # print(max_num, tops)
        for t in range(max_num-1, max(0, max_num-k-1), -1):
            assert t in tops


def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
